///
/// Copyright (C) 2015-2016, Dependable Systems Laboratory, EPFL
/// Copyright (C) 2015-2016, Cyberhaven
/// All rights reserved.
///
/// Licensed under the Cyberhaven Research License Agreement.
///

#ifndef S2E_PLUGINS_CGC_INTERFACE_H
#define S2E_PLUGINS_CGC_INTERFACE_H

#include <s2e/Plugin.h>
#include <s2e/Plugins/Core/BaseInstructions.h>
#include <s2e/Plugins/Coverage/BasicBlockCoverage.h>
#include <s2e/Plugins/Coverage/TranslationBlockCoverage.h>
#include <s2e/Plugins/ExecutionMonitors/CallSiteMonitor.h>
#include <s2e/Plugins/Models/StaticFunctionModels.h>
#include <s2e/Plugins/OSMonitors/Support/ModuleExecutionDetector.h>
#include <s2e/Plugins/Searchers/SeedSearcher.h>
#include <s2e/Synchronization.h>

#include "PovGenerationPolicy.h"

#include <unordered_set>
#include <vector>

namespace s2e {
namespace plugins {

enum TC_TYPES { TC_EXP1, TC_EXP2, TC_TEST, TC_CRASH, TC_TIMEOUT };
enum BB_TYPES { FUNC_ENTRY = 1, ICALL_TGT = 1 << 1, IJMP_TGT = 1 << 2, JIT_ENTRY = 1 << 3 };

namespace pov {
class PovGenerator;
}

class DecreeMonitor;
class ModuleExecutionDetector;

struct CBStats {
    /* Whether the CB called the random() syscall at some point */
    bool calledRandom;

    /* Stores program counters of branches that depend on a random value */
    llvm::DenseSet<uint64_t> randomBranchesPc;
};

class CGCInterface : public Plugin {
    S2E_PLUGIN

    typedef PovGenerationPolicy::TestCaseType TestCaseType;

private:
    friend class CGCInterfaceState;
    DecreeMonitor *m_monitor;
    ModuleExecutionDetector *m_detector;
    ProcessExecutionDetector *m_procDetector;
    pov::DecreePovGenerator *m_povGenerator;
    PovGenerationPolicy *m_exploitGenerator;
    coverage::BasicBlockCoverage *m_coverage;
    coverage::TranslationBlockCoverage *m_tbcoverage;
    ControlFlowGraph *m_cfg;
    CallSiteMonitor *m_csTracker;
    seeds::SeedSearcher *m_seedSearcher;
    models::StaticFunctionModels *m_models;

    typedef std::set<u_int64_t> AddressSet;
    typedef std::unordered_map<uint64_t, uint32_t> BBFlags;
    BBFlags m_bbFlags;

    uint64_t m_maxPovCount;
    bool m_recordAllPaths;
    bool m_recordConstraints;
    bool m_disableSendingExtraDataToDB; // data is used by fuzzer

    std::unordered_map<std::string, CBStats> m_cbStats;
    uint64_t m_cbStatsLastSent;
    unsigned m_cbStatsUpdateInterval;
    bool m_cbStatsChanged;

    uint64_t m_timeOfLastCoverageReport;
    uint64_t m_coverageTimeout;

    coverage::GlobalCoverage m_coveredTbs;
    coverage::ModuleTBs m_localCoveredTbs;

    typedef pov::PovOptions PovOptions;
    typedef pov::PovType PovType;

    bool sendCoveragePov(S2EExecutionState *state, TestCaseType tctype);

    void onStateKill(S2EExecutionState *state);

    void onRandom(S2EExecutionState *state, uint64_t pid, const std::vector<klee::ref<klee::Expr>> &);

    void onRandomInputFork(S2EExecutionState *state, const ModuleDescriptor &module);

    void onTimer();

    void onPovReady(S2EExecutionState *state, const PovOptions &opt, const std::string &recipeName,
                    const std::vector<std::string> &filePaths, TestCaseType tcType);

    void sendTestcase(S2EExecutionState *state, const std::string &xmlPovPath, const std::string &cPovPath,
                      TestCaseType tcType, const PovOptions &opt, const std::string &recipeName = "");

    void constraintsToJson(S2EExecutionState *state, std::stringstream &output);
    std::string constraintsToJsonFile(S2EExecutionState *state);

    void processIntermediateCoverage(uint64_t currentTime);

    bool updateCoverage(S2EExecutionState *state);

public:
    CGCInterface(S2E *s2e) : Plugin(s2e) {
    }
    void initialize();
}; // class CGCInterface

} // namespace plugins
} // namespace s2e

#endif // S2E_PLUGINS_CGC_INTERFACE_H
